#!/usr/bin/env drgn

from drgn.helpers.common import escape_ascii_string
from drgn.helpers.linux.cpumask import for_each_online_cpu
from drgn.helpers.linux.list import list_for_each_entry
from drgn.helpers.linux.percpu import per_cpu
from drgn import Object


def print_cfs_runq(runqueue: Object) -> None:
    cfs_root = hex(runqueue.cfs.tasks_timeline.address_of_().value_())
    print(" CFS RB_ROOT:", cfs_root)
    count = 0
    runq = runqueue.address_of_()
    cfs_rq = runqueue.cfs
    nr_running = cfs_rq.nr_running.value_()
    h_nr_running = cfs_rq.h_nr_running.value_()
    print(f'     nr_running: {nr_running}  h_nr_running: {h_nr_running}')
    for t in list_for_each_entry(
        "struct task_struct", runq.cfs_tasks.address_of_(), "se.group_node"):
        if t == runqueue.curr:
            continue
        count += 1
        print(
            " " * 4,
            '[{:3d}] PID: {:<6d} TASK: {}  COMMAND: "{}"'.format(
                t.prio.value_(),
                t.pid.value_(),
                hex(t),
                escape_ascii_string(t.comm.string_()),
            ),
        )
    if count == 0:
        print("     [no tasks queued]")

def print_job_runq(runqueue: Object) -> None:
    #job_root = hex(runqueue.job.tasks_timeline.address_of_().value_())
    #print("  JOB RB_ROOT:", job_root)
    count = 0
    runq = runqueue.address_of_()
    job_rq = runqueue.job

    nr_running = job_rq.nr_running.value_()
    h_nr_running = job_rq.h_nr_running.value_()
    print(f'     nr_running: {nr_running}  h_nr_running: {h_nr_running}')
    print("     job tasks:")
    for t in list_for_each_entry(
        "struct task_struct", runq.job_tasks.address_of_(), "job.group_node"):
        if t == runqueue.curr:
            continue
        count += 1
        print(
            " " * 4,
            '[{:3d}] PID: {:<6d} TASK: {}  COMMAND: "{}"'.format(
                t.prio.value_(),
                t.pid.value_(),
                hex(t),
                escape_ascii_string(t.comm.string_()),
            ),
        )
    if count == 0:
        print("     [no tasks queued]")


def main():
    for cpus in for_each_online_cpu(prog):
        runqueue = per_cpu(prog["runqueues"], cpus)
        comm = escape_ascii_string(runqueue.curr.comm.string_())
        pid = runqueue.curr.pid.value_()
        prio = runqueue.curr.prio.value_()
        print(f"CPU {cpus} RUNQUEUE: {runqueue.address_of_().value_():#x}")
        print(
            f"  curr:  PID: {pid:<6d}  TASK: {runqueue.curr.value_():#x}"
            f'  command: "{comm}"  prio: {prio}',
        )

        print_cfs_runq(runqueue)
        #print_job_runq(runqueue)
        


if __name__ == '__main__':
    main()
